package top.h1234.codesandbox;

import cn.hutool.core.util.ArrayUtil;
import com.github.dockerjava.api.DockerClient;
import com.github.dockerjava.api.async.ResultCallback;
import com.github.dockerjava.api.command.*;
import com.github.dockerjava.api.model.*;
import com.github.dockerjava.core.DockerClientBuilder;
import com.github.dockerjava.core.command.ExecStartResultCallback;
import org.springframework.stereotype.Component;
import org.springframework.util.StopWatch;
import top.h1234.codesandbox.model.ExecuteCodeRequest;
import top.h1234.codesandbox.model.ExecuteCodeResponse;
import top.h1234.codesandbox.model.ExecuteMessage;

import java.io.Closeable;
import java.io.File;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * Java代码沙箱 —— Docker实现
 */
@Component
public class JavaDockerCodeSandbox extends JavaCodeSandboxTemplate {

    @Override
    public ExecuteCodeResponse executeCode(ExecuteCodeRequest executeCodeRequest) {
        return super.executeCode(executeCodeRequest);
    }

    /**
     * 运行代码（覆盖模板中的方法）
     *
     * @param inputCaseList 输入用例列表
     * @param file          代码文件
     */
    @Override
    public List<ExecuteMessage> runFile(List<String> inputCaseList, File file) {
        // 返回的执行结果信息列表
        List<ExecuteMessage> executeMessageList = new ArrayList<>();

        // 用户代码父级文件路径
        String userCodeParentPath = file.getParentFile().getPath();
        // Docker客户端
        DockerClient dockerClient = DockerClientBuilder.getInstance().build();

        // 拉取镜像
        String imageName = "openjdk:8-alpine";
        pullImage(dockerClient, imageName);

        // 创建容器
        String containerId = createContainer(dockerClient, imageName, userCodeParentPath);

        // 运行容器
        dockerClient.startContainerCmd(containerId).exec();

        // 容器内运行代码
        for (String inputCase : inputCaseList) {
            String[] inputArgArr = inputCase.split(" ");
            // 执行命令（docker exec 容器名 java -cp /app Main 1 3）
            String[] cmdArray = ArrayUtil.append(new String[]{"java", "-cp", "/app", "Main"}, inputArgArr);
            // 创建执行命令
            ExecCreateCmd execCreateCmd = dockerClient.execCreateCmd(containerId);
            ExecCreateCmdResponse execCreateCmdResponse = execCreateCmd
                    .withCmd(cmdArray)
                    .withAttachStdin(true)
                    .withAttachStdout(true)
                    .withAttachStderr(true)
                    .exec();
            System.out.println("创建执行命令响应：" + execCreateCmdResponse);

            // 执行命令id
            String execId = execCreateCmdResponse.getId();

            // 执行已创建好的命令并获取消耗的内存、时间、执行的结果
            execAndGetResultMessage(dockerClient, containerId, execId, executeMessageList);
        }
        // 关闭并删除容器
        dockerClient.stopContainerCmd(containerId).exec();
        dockerClient.removeContainerCmd(containerId).withForce(true).exec();

        return executeMessageList;
    }


    /**
     * 执行已创建好的命令并获取消耗的内存、时间、执行的结果
     *
     * @param dockerClient       docker客户端
     * @param containerId        容器id
     * @param execId             执行命令id
     * @param executeMessageList 执行结果消息列表
     */
    public static void execAndGetResultMessage(DockerClient dockerClient, String containerId, String execId, List<ExecuteMessage> executeMessageList) {
        // 执行结果消息对象
        ExecuteMessage executeMessage = new ExecuteMessage();
        // 最大内存占用
        final long[] maxMemory = {0L};
        // 获取统计的数据（内存等）
        StatsCmd statsCmd = dockerClient.statsCmd(containerId);
        ResultCallback<Statistics> statisticsResultCallback = new ResultCallback<Statistics>() {
            // 当获取到统计信息时，这个方法会被调用
            @Override
            public void onNext(Statistics statistics) {
                // 获取最大内存的使用情况
                Long usage = statistics.getMemoryStats().getUsage();
                System.out.println("内存占用：" + usage + " B");
                maxMemory[0] = Math.max(usage, maxMemory[0]);
            }

            // 在获取统计信息的操作开始时被调用
            @Override
            public void onStart(Closeable closeable) {
            }

            // 在获取统计信息的过程中出现错误
            @Override
            public void onError(Throwable throwable) {
            }

            // 当获取统计信息的操作完成时（无论成功与否），这个方法都会被调用
            @Override
            public void onComplete() {
            }

            // 关闭回调的通用方法，用于释放与该回调相关的资源
            @Override
            public void close() {
            }
        };
        statsCmd.exec(statisticsResultCallback);

        // 程序执行耗时
        long time;
        // 是否超时标识
        final boolean[] timeout = {true};
        // 错误消息
        final String[] errMessage = {null};
        // 正常消息
        final String[] message = {null};
        ExecStartResultCallback execStartResultCallback = new ExecStartResultCallback() {
            @Override
            public void onNext(Frame frame) {
                // 获取流的类型
                StreamType streamType = frame.getStreamType();
                // 错误流
                if (StreamType.STDERR.equals(streamType)) {
                    errMessage[0] = new String(frame.getPayload());
                    System.out.print("【错误消息】" + errMessage[0]);
                } else {
                    message[0] = new String(frame.getPayload());
                    System.out.print("【正常消息】" + message[0]);
                }
                super.onNext(frame);
            }

            @Override
            public void onComplete() {
                // 程序正常结束，将超时标识改为false未超时
                timeout[0] = false;
                super.onComplete();
            }
        };

        // 执行开始（启动已创建的执行命令）
        try {
            // 开始计时
            StopWatch stopWatch = new StopWatch();
            stopWatch.start();

            dockerClient.execStartCmd(execId)
                    .exec(execStartResultCallback)
                    .awaitCompletion(5000L, TimeUnit.MILLISECONDS); // 最多运行执行5秒

            // 停止计时
            stopWatch.stop();
            // 停止统计内存
            statsCmd.close();
            if (timeout[0]) {
                throw new RuntimeException("超时运行");
            }
            time = stopWatch.getLastTaskTimeMillis();
            System.out.println("耗时：" + time + "毫秒");
        } catch (InterruptedException e) {
            System.out.println("程序执行异常");
            throw new RuntimeException(e);
        }

        // 封装结果
        executeMessage.setMessage(message[0]);
        executeMessage.setErrorMessage(errMessage[0]);
        executeMessage.setTime(time);
        executeMessage.setMemory(maxMemory[0]);
        executeMessageList.add(executeMessage);
    }

    /**
     * 创建容器，并将代码文件复制到容器内
     *
     * @param dockerClient       docker客户端
     * @param imageName          镜像名称
     * @param userCodeParentPath 用户代码存放的父级目录
     * @return 容器id
     */
    public static String createContainer(DockerClient dockerClient, String imageName, String userCodeParentPath) {
        CreateContainerCmd containerCmd = dockerClient.createContainerCmd(imageName);
        // DockerHost配置
        HostConfig hostConfig = new HostConfig();
        // 设置数据卷绑定（主机上的 userCodeParentPath 目录将被映射到容器内部的 /app 目录）
        hostConfig.setBinds(new Bind(userCodeParentPath, new Volume("/app")));
        // 设置内存最大使用量（字节，100M）
        hostConfig.withMemory(100 * 1000 * 1000L);
        hostConfig.withMemorySwap(0L);
        // 设置CPU最大使用数量
        hostConfig.withCpuCount(1L);

        CreateContainerResponse createContainerResponse = containerCmd
                .withHostConfig(hostConfig)
                .withNetworkDisabled(true) // 禁用网络
                .withReadonlyRootfs(true)          // 禁止往root中写文件
                .withAttachStdin(true)     // 标准输入
                .withAttachStdout(true)    // 标准输出
                .withAttachStderr(true)    // 标准错误输出
                .withTty(true)             // 开始终端交互（守护进程）
                .exec();

        // 获取刚刚创建的容器的容器id
        return createContainerResponse.getId();
    }

    /**
     * 拉取镜像
     *
     * @param dockerClient docker客户端
     * @param imageName    镜像名称
     */
    public static void pullImage(DockerClient dockerClient, String imageName) {
        PullImageCmd pullImageCmd = dockerClient.pullImageCmd(imageName);
        PullImageResultCallback pullImageResultCallback = new PullImageResultCallback() {
            @Override
            public void onNext(PullResponseItem item) {
                System.out.println("下载镜像：" + item.getStatus());
                super.onNext(item);
            }
        };
        try {
            pullImageCmd.exec(pullImageResultCallback).awaitCompletion();
        } catch (InterruptedException e) {
            System.out.println("下载镜像出错：" + e.getMessage());
            throw new RuntimeException(e);
        }
        System.out.println("下载 " + imageName + " 完成");
    }

}
